# 处理非线性stokes方程 删除非线性项
# Reference He's Book 5.3 稳态Navier-Stokes方程
from mpi4py import MPI
import numpy as np

import ufl
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, la, log, mesh
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    form,
    functionspace,
)
from petsc4py import PETSc
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block, NonlinearProblem, LinearProblem
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
from ufl import div, dx, inner, nabla_grad, dot, TestFunction, TrialFunction, TestFunctions, TrialFunctions, SpatialCoordinate
import ufl.functionspace
from dolfinx.nls.petsc import NewtonSolver


# Solution
def u_x(x):
    """
    x方向真实解
    """
    return x[0]**2 * x[1]**2 + np.exp(-x[1])


def u_y(x):
    """
    y方向真实解
    """
    return -2/3 * x[0] * x[1]**3 + 2 - np.pi * np.sin(np.pi * x[0])


def pressure(x):
    """
    压强的真实解
    """
    return -(2 - np.pi * np.sin(np.pi * x[0])) * np.cos(2 * np.pi * x[1])


def true_velocity(x):
    """
    获取真实速度
    return: List
    """
    return [u_x(x), u_y(x)]


def f_x(x_list):
    x, y = x_list  # 解包输入的x和y
    term1 = -2 * nu * x ** 2 - 2 * nu * y ** 2 - nu * ufl.exp(-y)
    term2 = ufl.pi ** 2 * ufl.cos(ufl.pi * x) * ufl.cos(2 * ufl.pi * y)
    term3 = 2 * x * y ** 2 * (x ** 2 * y ** 2 + ufl.exp(-y))
    term4 = (-2 * x * y ** 3 / 3 + 2 - ufl.pi * ufl.sin(ufl.pi * x)) * \
        (2 * x ** 2 * y - ufl.exp(-y))
    return term1 + term2


def f_y(x_list):
    x, y = x_list  # 解包输入的x和y
    term1 = 4 * nu * x * y - nu * ufl.pi ** 3 * ufl.sin(ufl.pi * x)
    term2 = 2 * ufl.pi * (2 - ufl.pi * ufl.sin(ufl.pi * x)
                          ) * ufl.sin(2 * ufl.pi * y)
    term3 = (x ** 2 * y ** 2 + ufl.exp(-y)) * \
        (-2 * y ** 3 / 3 - ufl.pi ** 2 * ufl.cos(ufl.pi * x))
    term4 = (-2 * x * y ** 3 / 3 + 2 - ufl.pi *
             ufl.sin(ufl.pi * x)) * (-2 * x * y ** 2)
    return term1 + term2


def f(x):
    return [f_x(x), f_y(x)]


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def get_norm():
    # 构造真解
    u_true = Function(Q)
    p_true = Function(S)
    u_true.interpolate(true_velocity)
    p_true.interpolate(pressure)

    u_error = Function(Q)
    u_h = sol.sub(0).collapse()
    p_h = sol.sub(1).collapse()
    u_error.x.array[:] = u_h.x.array - u_true.x.array
    p_error = Function(S)
    p_error.x.array[:] = p_h.x.array - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((ufl.grad(u_error)**2 + u_error**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p_h.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((ufl.grad(p_error)**2 + p_error**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)

    if MPI.COMM_WORLD.rank == 0:
        # 打开文件以追加模式写入
        with open('linear-stokes.txt', 'a') as file:
            file.write(f"N = {n_list[idx]}\n")
            file.write(
                f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.\n")
            file.write(
                f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.\n")
            file.write(
                f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.\n")


# Params
n_list = [
    [16, 4],
    [32, 8],
    [64, 16],
    [128, 32],
    [256, 64],
]
idx = int(input("请输入网格划分序号："))  # 0-4
nu = 1

msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, -0.25]), np.array(
        [1, 0])], n=n_list[idx], cell_type=CellType.triangle
)

# Function Space
P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
mel = mixed_element([P2, P1])  # 定义混合元空间 分别对应速度 压力
ME = functionspace(msh, mel)  # 实例化混合元空间

x = SpatialCoordinate(msh)  # 定义空间坐标


Q, _ = ME.sub(0).collapse()
S, _ = ME.sub(1).collapse()
u_bc = Function(Q)  # 真实速度函数空间
p_bc = Function(S)  # 真实压力函数空间
u_bc.interpolate(true_velocity)
p_bc.interpolate(pressure)

sol = Function(ME)  # 定义解函数
sol.x.array[:] = np.random.randn() - 0.5  # 不初始化为0 无法收敛

# type u, p =  <class 'ufl.tensors.ListTensor'>, <class 'ufl.indexed.Indexed'>
u, p = ufl.split(sol)  # trial: u, p; test: v, q
v, q = TestFunctions(ME)  # Test Functions

# Weak Formulation
F = 2 * nu * inner(D(u), D(v)) * dx
F += inner(div(u), q) * dx
F -= inner(div(v), p) * dx
# F += inner(dot(u, nabla_grad(u)), v) * dx  # 非线性项
# test
f_expr = ufl.as_vector([f_x(x), f_y(x)])  # 调整f为ufl格式
F -= inner(f_expr, v) * dx  # 右端项


# BCs
fdim = msh.topology.dim - 1
facets = mesh.locate_entities_boundary(
    msh, fdim, lambda x: np.isclose(x[0], 1.0) | np.isclose(
        x[0], 0.0) | np.isclose(x[1], -0.25) | np.isclose(x[1], 0.0)
)
# 自由度的空间使用(V0, Q) 边界条件设置为V0
v_dofs = fem.locate_dofs_topological((ME.sub(0), Q), fdim, facets)
# 当传入value是插值函数 会自动导入函数空间 不需要设置空间
bc_vel = dirichletbc(u_bc, v_dofs, ME.sub(0))  # 速度边界条件

# 设置压力


def marker(x): return np.isclose(x[0], 0.0) & np.isclose(x[1], -0.25)


p_dofs = fem.locate_dofs_geometrical((ME.sub(1), S), marker)
bc_pressure = dirichletbc(p_bc, p_dofs, ME.sub(1))


# Create nonlinear problem and Newton solver
problem = NonlinearProblem(F, sol, bcs=[bc_vel, bc_pressure])
solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.convergence_criterion = "incremental"
solver.rtol = 1e-8
solver.max_it = 500
solver.report = True

ksp = solver.krylov_solver
opts = PETSc.Options()
option_prefix = ksp.getOptionsPrefix()
opts[f"{option_prefix}ksp_type"] = "preonly"  # 用cg就一直不行
opts[f"{option_prefix}pc_type"] = "lu"
opts[f"{option_prefix}pc_factor_mat_solver_type"] = "superlu_dist"
# 唯一能算的方法
# opts[f"{option_prefix}pc_type"] = "gamg"  # gamg不需要设置pc_factor_mat_solver_type
ksp.setFromOptions()

log.set_log_level(log.LogLevel.INFO)
n, converged = solver.solve(sol)
assert (converged)
print(f"Number of interations: {n:d}")

# 计算误差
get_norm()
